{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "75b09255-c46f-42b2-9e92-5e64861efa43",
   "metadata": {},
   "source": [
    "# Qwen3微调实战：医疗R1推理风格聊天\n",
    "\n",
    "[![](https://raw.githubusercontent.com/SwanHubX/assets/main/badge1.svg)](https://swanlab.cn/@ZeyiLin/qwen3-sft-medical/overview)\n",
    "\n",
    "- **Github**: [Qwen3-Medical-SFT](https://github.com/Zeyi-Lin/Qwen3-Medical-SFT)\n",
    "- **基础模型**：[Qwen3-1.7B](https://modelscope.cn/models/Qwen/Qwen3-1.7B/summary)\n",
    "- **微调后模型**：[Qwen3-1.7b-Medical-R1-sft](https://modelscope.cn/models/testUser/Qwen3-1.7b-Medical-R1-sft/summary)\n",
    "- **数据集**：[delicate_medical_r1_data](https://modelscope.cn/datasets/krisfu/delicate_medical_r1_data)\n",
    "- **SwanLab**：[qwen3-sft-medical](https://swanlab.cn/@ZeyiLin/qwen3-sft-medical/runs/agps0dkifth5l1xytcdyk/chart)\n",
    "- **微调方式**：全参数微调、LoRA微调\n",
    "- **推理风格**：R1推理风格\n",
    "- **算力要求**：\n",
    "  - **全参数微调**：32GB显存\n",
    "  - **LoRA微调**：28GB显存\n",
    "- **图文教程**：[Qwen3大模型微调入门实战（完整代码）](https://zhuanlan.zhihu.com/p/1903848838214705484)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21b4fa4a-03a4-4bdb-b593-b50d2ebbb931",
   "metadata": {},
   "source": [
    "## 1. 安装环境"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42630f8a-e3d2-4972-9a44-a36aec69675f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c73ce7d-54a9-4928-9b27-c58c68c090d7",
   "metadata": {},
   "source": [
    "## 2. 下载数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7c0912c-ab6e-4d4d-95b5-ba9e11687d2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python data.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ac44ba5-45d2-4e43-a906-2e4421b72114",
   "metadata": {},
   "source": [
    "## 3. 登录SwanLab\n",
    "1. 前往[swanlab](https://swanlab.cn/space/~/settings)复制你的API Key，粘贴到下面的代码中\n",
    "2. 如果你不希望将登录信息保存到该计算机中，可将`save=True`去掉（每次运行训练需要重新执行下面的代码块）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3d5cbc21-ae61-40fd-ae38-843b774fe82c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                                                    "
     ]
    }
   ],
   "source": [
    "import swanlab\n",
    "\n",
    "swanlab.login(api_key=\"[你的API Key]\", save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f43e7f4-976b-46ec-ab63-febeac40a3ce",
   "metadata": {},
   "source": [
    "## 4. 开启全参数微调"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c451e2e-8542-404c-9eab-8fd610bf8e12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import torch\n",
    "from datasets import Dataset\n",
    "from modelscope import snapshot_download, AutoTokenizer\n",
    "from transformers import AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForSeq2Seq\n",
    "import os\n",
    "import swanlab\n",
    "\n",
    "os.environ[\"SWANLAB_PROJECT\"]=\"qwen3-sft-medical\"\n",
    "PROMPT = \"你是一个医学专家，你需要根据用户的问题，给出带有思考的回答。\"\n",
    "MAX_LENGTH = 2048\n",
    "\n",
    "swanlab.config.update({\n",
    "    \"model\": \"Qwen/Qwen3-1.7B\",\n",
    "    \"prompt\": PROMPT,\n",
    "    \"data_max_length\": MAX_LENGTH,\n",
    "    })\n",
    "\n",
    "def dataset_jsonl_transfer(origin_path, new_path):\n",
    "    \"\"\"\n",
    "    将原始数据集转换为大模型微调所需数据格式的新数据集\n",
    "    \"\"\"\n",
    "    messages = []\n",
    "\n",
    "    # 读取旧的JSONL文件\n",
    "    with open(origin_path, \"r\") as file:\n",
    "        for line in file:\n",
    "            # 解析每一行的json数据\n",
    "            data = json.loads(line)\n",
    "            input = data[\"question\"]\n",
    "            output = f\"<think>{data[\"think\"]}</think> \\n {data[\"answer\"]}\"\n",
    "            message = {\n",
    "                \"instruction\": PROMPT,\n",
    "                \"input\": f\"{input}\",\n",
    "                \"output\": output,\n",
    "            }\n",
    "            messages.append(message)\n",
    "\n",
    "    # 保存重构后的JSONL文件\n",
    "    with open(new_path, \"w\", encoding=\"utf-8\") as file:\n",
    "        for message in messages:\n",
    "            file.write(json.dumps(message, ensure_ascii=False) + \"\\n\")\n",
    "\n",
    "\n",
    "def process_func(example):\n",
    "    \"\"\"\n",
    "    将数据集进行预处理\n",
    "    \"\"\" \n",
    "    input_ids, attention_mask, labels = [], [], []\n",
    "    instruction = tokenizer(\n",
    "        f\"<|im_start|>system\\n{PROMPT}<|im_end|>\\n<|im_start|>user\\n{example['input']}<|im_end|>\\n<|im_start|>assistant\\n\",\n",
    "        add_special_tokens=False,\n",
    "    )\n",
    "    response = tokenizer(f\"{example['output']}\", add_special_tokens=False)\n",
    "    input_ids = instruction[\"input_ids\"] + response[\"input_ids\"] + [tokenizer.pad_token_id]\n",
    "    attention_mask = (\n",
    "        instruction[\"attention_mask\"] + response[\"attention_mask\"] + [1]\n",
    "    )\n",
    "    labels = [-100] * len(instruction[\"input_ids\"]) + response[\"input_ids\"] + [tokenizer.pad_token_id]\n",
    "    if len(input_ids) > MAX_LENGTH:  # 做一个截断\n",
    "        input_ids = input_ids[:MAX_LENGTH]\n",
    "        attention_mask = attention_mask[:MAX_LENGTH]\n",
    "        labels = labels[:MAX_LENGTH]\n",
    "    return {\"input_ids\": input_ids, \"attention_mask\": attention_mask, \"labels\": labels}   \n",
    "\n",
    "\n",
    "def predict(messages, model, tokenizer):\n",
    "    device = \"cuda\"\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        messages,\n",
    "        tokenize=False,\n",
    "        add_generation_prompt=True\n",
    "    )\n",
    "    model_inputs = tokenizer([text], return_tensors=\"pt\").to(device)\n",
    "\n",
    "    generated_ids = model.generate(\n",
    "        model_inputs.input_ids,\n",
    "        max_new_tokens=MAX_LENGTH,\n",
    "    )\n",
    "    generated_ids = [\n",
    "        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)\n",
    "    ]\n",
    "\n",
    "    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
    "\n",
    "    return response\n",
    "\n",
    "# 在modelscope上下载Qwen模型到本地目录下\n",
    "model_dir = snapshot_download(\"Qwen/Qwen3-1.7B\", cache_dir=\"./\", revision=\"master\")\n",
    "\n",
    "# Transformers加载模型权重\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"./Qwen/Qwen3-1.7B\", use_fast=False, trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(\"./Qwen/Qwen3-1.7B\", device_map=\"auto\", torch_dtype=torch.bfloat16)\n",
    "model.enable_input_require_grads()  # 开启梯度检查点时，要执行该方法\n",
    "\n",
    "# 加载、处理数据集和测试集\n",
    "train_dataset_path = \"train.jsonl\"\n",
    "test_dataset_path = \"val.jsonl\"\n",
    "\n",
    "train_jsonl_new_path = \"train_format.jsonl\"\n",
    "test_jsonl_new_path = \"val_format.jsonl\"\n",
    "\n",
    "if not os.path.exists(train_jsonl_new_path):\n",
    "    dataset_jsonl_transfer(train_dataset_path, train_jsonl_new_path)\n",
    "if not os.path.exists(test_jsonl_new_path):\n",
    "    dataset_jsonl_transfer(test_dataset_path, test_jsonl_new_path)\n",
    "\n",
    "# 得到训练集\n",
    "train_df = pd.read_json(train_jsonl_new_path, lines=True)\n",
    "train_ds = Dataset.from_pandas(train_df)\n",
    "train_dataset = train_ds.map(process_func, remove_columns=train_ds.column_names)\n",
    "\n",
    "# 得到验证集\n",
    "eval_df = pd.read_json(test_jsonl_new_path, lines=True)\n",
    "eval_ds = Dataset.from_pandas(eval_df)\n",
    "eval_dataset = eval_ds.map(process_func, remove_columns=eval_ds.column_names)\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=\"./output/Qwen3-1.7B\",\n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    gradient_accumulation_steps=4,\n",
    "    eval_strategy=\"steps\",\n",
    "    eval_steps=100,\n",
    "    logging_steps=10,\n",
    "    num_train_epochs=2,\n",
    "    save_steps=400,\n",
    "    learning_rate=1e-4,\n",
    "    save_on_each_node=True,\n",
    "    gradient_checkpointing=True,\n",
    "    report_to=\"swanlab\",\n",
    "    run_name=\"qwen3-1.7B\",\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    "    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "\n",
    "# 用测试集的前3条，主观看模型\n",
    "test_df = pd.read_json(test_jsonl_new_path, lines=True)[:3]\n",
    "\n",
    "test_text_list = []\n",
    "\n",
    "for index, row in test_df.iterrows():\n",
    "    instruction = row['instruction']\n",
    "    input_value = row['input']\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": f\"{instruction}\"},\n",
    "        {\"role\": \"user\", \"content\": f\"{input_value}\"}\n",
    "    ]\n",
    "\n",
    "    response = predict(messages, model, tokenizer)\n",
    "\n",
    "    response_text = f\"\"\"\n",
    "    Question: {input_value}\n",
    "\n",
    "    LLM:{response}\n",
    "    \"\"\"\n",
    "    \n",
    "    test_text_list.append(swanlab.Text(response_text))\n",
    "    print(response_text)\n",
    "\n",
    "swanlab.log({\"Prediction\": test_text_list})\n",
    "\n",
    "swanlab.finish()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
